import cv2
import numpy as np
# from capture import USBCamera

# def saveFrame():
# 	capture = USBCamera("/dev/video0")
# 	frame = capture.getFrame()
# 	cv2.imwrite('capture.png', frame)


# import numpy as np
# import glob
# size = [480,400]
# def get_K_and_D(checkerboard, imgsPath):
#     CHECKERBOARD = checkerboard
#     subpix_criteria = (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 0.1)
#     calibration_flags = cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC+cv2.fisheye.CALIB_CHECK_COND+cv2.fisheye.CALIB_FIX_SKEW
#     objp = np.zeros((1, CHECKERBOARD[0]*CHECKERBOARD[1], 3), np.float32)
#     objp[0,:,:2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2)
#     _img_shape = None
#     objpoints = [] 
#     imgpoints = [] 
#     images = glob.glob(imgsPath + '/*.png')
#     for fname in images:
#         img = cv2.imread(fname)
# 		img = cv2.resize(img, size)
#         if _img_shape == None:
#             _img_shape = img.shape[:2]
#         else:
#             assert _img_shape == img.shape[:2], "All images must share the same size."
        
#         gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
#         ret, corners = cv2.findChessboardCorners(gray, CHECKERBOARD,cv2.CALIB_CB_ADAPTIVE_THRESH+cv2.CALIB_CB_FAST_CHECK+cv2.CALIB_CB_NORMALIZE_IMAGE)
#         if ret == True:
#             objpoints.append(objp)
#             cv2.cornerSubPix(gray,corners,(3,3),(-1,-1),subpix_criteria)
#             imgpoints.append(corners)
#     N_OK = len(objpoints)
#     K = np.zeros((3, 3))
#     D = np.zeros((4, 1))
#     rvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
#     tvecs = [np.zeros((1, 1, 3), dtype=np.float64) for i in range(N_OK)]
#     rms, _, _, _, _ = \
#     cv2.fisheye.calibrate(
#                           objpoints,
#                           imgpoints,
#                           gray.shape[::-1],
#                           K,
#                           D,
#                           rvecs,
#                           tvecs,
#                           calibration_flags,
#                           (cv2.TERM_CRITERIA_EPS+cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6)
#                           )
#     DIM = _img_shape[::-1]
#     print("Found " + str(N_OK) + " valid images for calibration")
#     print("DIM=" + str(_img_shape[::-1]))
#     print("K=np.array(" + str(K.tolist()) + ")")
#     print("D=np.array(" + str(D.tolist()) + ")")
#     return DIM, K, D


# def OpenCVFisheyeImageUndistortion(img):
	# cv::Mat img, undistortImg
	# cv::Matx33d K, P
	# cv::Vec4d D
	# cv::Mat mapX, mapY
	
	# K(0, 0) = 348.52; K(0, 1) = 0; K(0, 2) = 640.19
	# K(1, 0) = 0; K(1, 1) = 348.52; K(1, 2) = 358.56
	# K(2, 0) = 0; K(2, 1) = 0; K(2, 2) = 1
    
	# D(0) = 0.066258
	# D(1) = 0.039769
	# D(2) = -0.026906
	# D(3) = 0.003342
 
	# P = K
	# P(0, 0) /= 1.5
	# P(1, 1) /= 1.5

    # K = [[348.52, 0,      640.19], 
    #      [0,      348.52, 358.56],
    #      [0,      0,      1]]
    # D = [0.066258, 0.039769, -0.026906, 0.003342]
	# cv2.initUndistortRectifyMap(K, D, cv::Matx33d::eye(), P, cv::Size(img.cols, img.rows), CV_16SC2, mapX, mapY)
 
	# cv::remap(img, undistortImg, mapX, mapY, CV_INTER_LINEAR)
	# cv::imshow("src", img)
	# cv::imshow("corrected", undistortImg)



# DIM = (480, 400)
# K = np.array([[122.87238744667862, 0.0, 220.28343964853087], 
# 			  [0.0, 120.60935347262523, 203.74634232157828], 
# 			[0.0, 0.0, 1.0]])
# D = np.array([[0.06701341388742599], [0.1906721369775406], [-0.2568929474339438], [0.09063021944980182]])
DIM = (640,480)
K = np.array([[182.2034195269173, 0.0, 248.72241643510006], 
			  [0.0, 200.88841301751464, 195.87207910703572], 
			[0.0, 0.0, 1.0]])
D = np.array([[0.12236488620142214], [-0.053849293887633265], [0.03508527000849206], [-0.011661429257952398]])

def undistort(img, DIM=DIM, K=K, D=D):
    # img = cv2.resize(img, DIM)
    map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), K, DIM,cv2.CV_16SC2)
    undistorted_img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT)
    return undistorted_img


def get_img_middle(img, size=[640,480]):
	top = int((img.shape[0] - size[1]) / 2)
	bottom = top + size[1]
	left = int((img.shape[1] - size[0]) / 2)
	right = left + size[0]
	return img[top:bottom, left:right]




if __name__ == "__main__":
	capture = USBCamera("/dev/video0")
	frame = capture.getFrame()
	frame = cv2.resize(frame, DIM)
	map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, D, np.eye(3), K, DIM,cv2.CV_16SC2)
	undistorted_img = cv2.remap(img, map1, map2, interpolation=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT)
	cv2.imwrite('unfisheyeImage.png', undistorted_img)